//#pragma once
#ifndef __MALRTF_H__
#define __MALRTF_H__

#include "stdafx.h"



class RTAnalyzer: public refable_impl_t
{
#define RT_SECT_SIZE			0x400
	enum ParserStatus
	{
		PARSE_MAIN, 
		PARSE_CONTROL_, 
		PARSE_CONTROL_WORD, 
		PARSE_CONTROL_SYMBOL, 
		PARSE_CONTROL_WORD_PARAM, 
		PARSE_INTERPRET_CONTROLWORD 
	};
public:
	RTAnalyzer():ctrl_word_count(0),
		ctrl_symbol_cout(0),
		context_char_count(0),
		char_total_count(0),
		invalid_symbol_count(0)
	{}
public:
	HRESULT Parse(IKBStream* pstm)
	{
		RASSERT(pstm , E_FAIL);

		UTIL::sentry<UINT8*> pSect(new UINT8[RT_SECT_SIZE]);
		RASSERT(pSect , E_FAIL);
		RFAILED(pstm->Seek(0 , NULL,SPBegin));

		DWORD dw = 0;
		pstm->GetSize(&dw , NULL);

		m_parse_status =  PARSE_MAIN;

		UINT32 _file_size = dw;
		UINT32 _remain_bytes = _file_size;
		UINT32 _read_bytes = 0;
		while (_remain_bytes > 0)
		{
			_read_bytes = _remain_bytes > RT_SECT_SIZE ? RT_SECT_SIZE : _remain_bytes;
			RFAILED(pstm->Read(pSect , _read_bytes , NULL));
			_remain_bytes -= _read_bytes;
			RFAILED(ParseSection2(pSect , _read_bytes));
		}


		return S_OK;
	}
private:	
	HRESULT ParseSection(UINT8* pSect , UINT32 size)
	{
		//RASSERT(pBegin && pEnd , FALSE );
		UINT8* _bp = pSect;
		UINT8* _ep = pSect + size;

		while (_bp < _ep)
		{
			//char_total_count++;
			switch(m_parse_status) {
			case PARSE_MAIN: 
				switch(*_bp) 
				{
				case '{':
					_bp++;
					break;
				case '}':
					{
						_bp++;
					}
					break;
				case '\\':
					m_parse_status = PARSE_CONTROL_;
					_bp++;
					break;
				default:
					{
						if (isprint(*_bp)) char_total_count++;
						else invalid_symbol_count++;
						_bp++;
						context_char_count++;
					}	
					break;
				}
				break;
			case PARSE_CONTROL_:
				if(isalpha(*_bp))  {
					m_parse_status = PARSE_CONTROL_WORD;
				}
				else
					m_parse_status = PARSE_CONTROL_SYMBOL;
				break;
			case PARSE_CONTROL_SYMBOL:
				/*if (IsRtfSymbol(*_pBegin)) ctrl_symbol_cout++;
				else invalid_symbol_count++;*/
				_bp++;
				m_parse_status = PARSE_MAIN;
				break;
			case PARSE_CONTROL_WORD:
				//ctrl_word_count++;
				if(isalpha(*_bp))
					_bp++;
				else {
					if(isspace(*_bp)) {
						_bp++;
						m_parse_status = PARSE_INTERPRET_CONTROLWORD;
					}
					else if (isdigit(*_bp)) {
						m_parse_status = PARSE_CONTROL_WORD_PARAM;
					}
					else if(*_bp == '-') {
						_bp++;
						m_parse_status = PARSE_CONTROL_WORD_PARAM;
					}
					else {
						m_parse_status = PARSE_INTERPRET_CONTROLWORD;
					}
				}
				break;
			case PARSE_CONTROL_WORD_PARAM:
				if(isdigit(*_bp)) {
					_bp++;
				}
				else if(isalpha(*_bp)) {
					_bp++;
				}
				else {
					m_parse_status = PARSE_INTERPRET_CONTROLWORD;
				}
				break;
			case PARSE_INTERPRET_CONTROLWORD:
				{
					m_parse_status = PARSE_MAIN;
				}
				break;
			}

		}


		return S_OK;
	}
	HRESULT ParseSection2(UINT8* pSect , UINT32 size)
	{
		UINT8* pbeing = pSect;
		UINT8* pEnd = pSect + size;
		while (pbeing < pEnd)
		{
			if (!isprint(*pbeing)) invalid_symbol_count++;
			char_total_count++;
			pbeing++;


		}
		return S_OK;
	}
	inline BOOL IsRtfSymbol(UINT8 ch)
	{
		if (ch == '\'' || ch == '-' || ch == '*'
			|| ch == ':' || ch == '\\' || ch == '_'
			|| ch == '{' || ch == '|' || ch == '}'|| ch == '~') return TRUE;
		else return FALSE;
	}
public:
	UINT64 ctrl_word_count;
	UINT64 ctrl_symbol_cout;
	UINT64 invalid_symbol_count;
	UINT64 context_char_count;
	UINT64 char_total_count;
private:
	ParserStatus m_parse_status;
};


class MalRTFDetector :public IRXDetectNode, CUnknownImp
{
public:
	MalRTFDetector(void){};
	virtual ~MalRTFDetector(void){};
public:
	UNKNOWN_IMP1_(IRXDetectNode);
	HRESULT init_class( void * ,void * )
	{
		return S_OK;
	}
	STDMETHOD(Detect)( IUnknown *punkArchive, LPDETECTR lpResult )
	{
		optr<IRXArchive> oArchive( punkArchive );
		if( !oArchive ) return E_INVALIDARG;

		optr<IKBStream> oStream;
		RFAILED( oArchive->GetStreamObject( oStream.pp() ) );

		optr<IRXDTDataBox> box;
		RFAILED( ez::GetAuxDataBox( box, oStream ) );
		if( SUCCEEDED(box->DTFCFind( STR_SIZE_PAIR("|Feature.Avoid.") ) ) )
			return E_FAIL;

		RTAnalyzer _rtf_analyzer;
		RFAILED(_rtf_analyzer.Parse(oStream));

		DOUBLE rat = (DOUBLE)_rtf_analyzer.invalid_symbol_count/(DOUBLE)(_rtf_analyzer.char_total_count - _rtf_analyzer.invalid_symbol_count);
		//printf("char count : %lld		,invalid symbol count:%lld , rat : %f\n" ,_rtf_analyzer.char_total_count,_rtf_analyzer.invalid_symbol_count,rat);

		if (rat > 0.07) lpResult->Class = nam::kMalware;

		return S_OK;
	}

	STDMETHOD_(TRESULT,Treat)( IUnknown *punkArchive, LPDETECTR lpResult ) 
	{
		return SUCCEEDED(ez::RemoveByArchiveOrStream(punkArchive)) ? TR_S_TREAT : TR_E_DELETE;
	}

	STDMETHOD(Reset)()
	{
		return S_OK;
	}

};



#endif

